
#include "../include/Process.h"

void Process::ProcessNonlinear(){
    string fileType = ".txt";
    //! --------------- 1 Variable Assignment------------
    std::cout << "----1.Variable Assignment----" << std::endl;
    // Time--------------------------------------------
    auto t0     = Parameter::t0;;
    auto t_max  = Parameter::t_max;
    auto dt     = Parameter::dt;
    int  nt     = Parameter::nt;
    int fft_num = Parameter::fft_num;

    // Radial------------------------------------------
    auto r0     = Parameter::r0;
    auto dr_min = Parameter::dr_min;
    auto dr_max = Parameter::dr_max;
    int  nr     = Parameter::nr;
    int r_step_type = input::step_type;
    
    // Longitudinal------------------------------------
    auto z0 = Parameter::z0;
    auto dz = Parameter::dz;
    int  nz = Parameter::nz;

    // input laser1------------------------------------
    auto Phase1 = Parameter::phase1; //cep
    auto Phase2 = Parameter::phase2; //cep
    auto Phase0 = Parameter::phase0; //Relative Phase
    auto FWHM   = Parameter::fwhm;

	int  index_omega1 = Parameter::index_omega1; 
    auto Wavelength1  = Parameter::Wavelength1;
    auto Waist1       = Parameter::Waist1;
    auto Focus1       = Parameter::Focus1;
    auto E_peak1      = Parameter::E_peak1;
    
    // input laser2--------------------------------------
    int  index_omega2 = Parameter::index_omega2;
    auto Wavelength2  = Parameter::Wavelength2;
    auto Waist2       = Parameter::Waist2;
    auto Focus2       = Parameter::Focus2;
    auto E_peak2      = Parameter::E_peak2;

    // nonlinear parameter-----------------------------------
    auto Kerr = Parameter::Kerr;
    auto Ng   = Parameter::Ng;

    //---------------------------------
    int nr_out = Parameter::nr_out;
   
    cout<< "nr_out = " << nr_out <<endl;
    cout<< "nr = " << nr <<endl;
    cout<< "nz = " << nz <<endl;

    //!-----------------"Class Initialization"--------------------
    // std::cout << "----2.Class Initialization----" << std::endl;
	
    NonlinearResponse nonlinear(nt, dt, Kerr, Ng);
    CrankNicholson CN(nr);
    Step step(r_step_type, dr_min, dr_max, nr);
    
    MatrixAB matrixab(nr, fft_num, dz); 
    GaussianDual pulse(z0, FWHM, t_max, Phase0,
                        Waist1,Focus1,E_peak1,Wavelength1,Phase1,
                        Waist2,Focus2,E_peak2,Wavelength2,Phase2);
    

    //! ----------------------3 1D_Array-------------------------
    // std::cout << "----3. 1D_Array----" << std::endl;
    // Grid Generation----------------------------------------
    double *r_step_array = new double[nr]; // StepArray   
    double *R_grid = new double[nr];       // RadialGrid   
    double *T_grid = new double[nt];       // TimeGrid     
    double *Z_grid = new double[nz];       // LongitudinalGid 

    // NonliearResponse-----------------------------------------
    double *ionization_rate = new double[nt];
    double *electron_density= new double[nt];
    double *nonlinear_J     = new double[nt];
	double *nonlinear_P     = new double[nt]; 

    // Coefficient Matrix---------------------------------------

    double  *A0 = new double[nr];   
    double  *B0 = new double[nr];   
    double  *A2 = new double[nr];   
    double  *B2 = new double[nr];
    
    complex** A1 = new complex*[fft_num];
    complex** B1 = new complex*[fft_num];
    for(int i = 0; i < fft_num; i++){
        A1[i] = new complex[nr];
        B1[i] = new complex[nr];
    }

    // alfa，frequency_array-----------------------------------------
    double  *frequency_array = new double[nt]; 
    
    complex *fft_outc_E       = new complex[nt];     //E_fft
	complex *fft_outc_J       = new complex[nt];     //J_fft
    complex *fft_outc_P       = new complex[nt];     //P_fft
    // 1D_Array, Assignment----------------------------------------
    step(r_step_array);
    Grid::Grid_step(R_grid, nr, r0, r_step_array); // R_grid
    Grid::Grid_1d(T_grid, nt, dt, t0);             // T_grid
    Grid::Grid_1d(Z_grid, nz, dz, z0);             // Z_grid
    Fourier::fftfreq(nt,frequency_array, dt);              // frequency_array
    
    // A1,B1, Assignment----------------------------------------
    matrixab(frequency_array, 
            R_grid, r_step_array,
            A0, B0, A1, B1, A2, B2);   //Coefficient Matrix

	double d_omega = 2. * PI * fabs(frequency_array[1]-frequency_array[0]);
    int  index_THz = round((input::THz_max * Constants::omega1THz)/(d_omega));
    // cout <<"index_THz = "<<index_THz<<endl;

    //! ------------------ 4.  2D_Array----------------------------------------
    // std::cout << "----4.2D_Array----" << std::endl;
    Grid::Grid_2d *grid_2d = new Grid::Grid_2d(nr, nt); // 结束时释放
    // E_n0 initialization------------------------------------
    Grid::ini_zeros(nr, nt,       grid_2d->E_n0);
    Grid::ini_zeros(fft_num, nr,  grid_2d->E_n0fft);
    Grid::ini_zeros(fft_num, nr,  grid_2d->G_n0fft);
    Grid::ini_zeros(fft_num, nr,  grid_2d->G_n_1fft);
    Grid::ini_zeros(nr, fft_num,  grid_2d->E_n1fft);
    pulse(nr, nt, R_grid, T_grid, grid_2d->E_n0);

//!------------------5. Absorption Function---------------------------------
    //The thickness of the absorbing layer
    int ab_len = 100;       
    double ab_array[ab_len];
    Absord ab_function;
    ab_function(ab_array, ab_len);

//!------------------6. Output Folder----------------------------------------
    // cout << "----6.output_folder----" << endl;
    // grid_out-----------------------------------------------------------
    Fileout::fileout_grid(Parameter::filename_grid, frequency_array, T_grid, 
                          R_grid, Z_grid, ab_array, nt, nr, nz, ab_len);
    Parameter::fileout_parameter();
    Fileout::output_Erz_accuracy(Parameter::filename_Erz_accuracy, 
                          pulse, R_grid, Z_grid, nr_out, nz);

    //!----------------7 Release Space before loop-------------------------------------
    // std::cout << "7 Release Space before loop" << std::endl;
    delete[] Z_grid;
    delete[] T_grid;
    delete[] R_grid;
    delete[] r_step_array;

    //!-----------------8 into outer loop nz--------------------------------------------
    cout << "----8.into outer loop nz----" << endl;

    double intensity_array[nr_out];
    double intensityTHz_array[nr_out];
    double plasmaDensity_array[nr_out];
    double phase1_array[nr_out];
    double phase2_array[nr_out];

    for (int num_z = 0; num_z < nz; num_z++)
    {
        std::cout << "nz=" << num_z << std::endl;
        //! fft----------G_n0fft--E_n0fft -------------------
        int num_ab = 0;  

        for (int num_r = 0; num_r < nr; num_r++)
        {
            num_ab = (num_r < nr - ab_len)? 0 : num_r - nr + ab_len;
            nonlinear(ionization_rate, electron_density,
                      nonlinear_J, nonlinear_P,
                      grid_2d->E_n0[num_r]);

            Fourier::fft_r2c(nt, grid_2d->E_n0[num_r], fft_outc_E);
			Fourier::fft_r2c(nt, nonlinear_J, fft_outc_J);
            Fourier::fft_r2c(nt, nonlinear_P, fft_outc_P);


            complex iwJ = 0;
            complex w2P = 0;

            double intensity_all = 0;
			double intensity_THz = 0;
            double phase1_w  = 0;
            double phase2_w  = 0;
			double plasmaDensity = 0;

            for (int num_t = 0; num_t < fft_num; num_t++)
            { 
                // 将第num_r行的nt个数据变换后的数据，放入fft第num_r列
				double ab_i    = ab_array[num_ab];
                double omega_i = 2. * PI * frequency_array[num_t];

                grid_2d->E_n0fft[num_t][num_r] = fft_outc_E[num_t];
                grid_2d->E_n0fft[num_t][num_r] *= (num_r < nr - ab_len)? 1.0:ab_i;//吸收边

				iwJ = fft_outc_J[num_t] * omega_i * Constants::Ii;
				w2P =-fft_outc_P[num_t] * omega_i * omega_i;
                complex Gn0 = dz * (iwJ + w2P);
                complex Gn_1 = (num_z == 0)? Gn0 : grid_2d->G_n_1fft[num_t][num_r];

				// grid_2d->G_n0fft[num_t][num_r] = 2.*Gn0 ;//Adams-Bashforth差分格式
				grid_2d->G_n0fft[num_t][num_r] = 3.*Gn0 - Gn_1;//Adams-Bashforth差分格式
                grid_2d->G_n0fft[num_t][num_r] *= (num_r < nr - ab_len) ? 1.0: ab_i;//吸收边
                grid_2d->G_n_1fft[num_t][num_r] = Gn0; //将当前层的非线性存储到nz_1中，作为下一层的上一层数据

                intensity_all += fabs(grid_2d->E_n0fft[num_t][num_r]);
				intensity_THz = (num_t <= index_THz)? intensity_all: intensity_THz;
            }


            intensity_all = intensity_all/((nt*1.)/2.);
			intensity_THz = intensity_THz /((nt*1.)/2.);
            phase1_w  = Util::Get_phase(grid_2d->E_n0fft[index_omega1][num_r]);
			plasmaDensity = electron_density[nt-1];

            intensity_array[num_r]     = intensity_all;
			intensityTHz_array[num_r]  = intensity_THz;
            plasmaDensity_array[num_r] = plasmaDensity;
            phase1_array[num_r]        = phase1_w;
            
            if(E_peak2!=0){
                phase2_w  = Util::Get_phase(grid_2d->E_n0fft[index_omega2][num_r]);        
                phase2_array[num_r]      = phase2_w;
            }
            if(num_r == 0){
                Fileout::output_Ewzc(Parameter::THzSpectrumyOut, nt, fft_outc_E);
                std::cout<< "intensity: "<< intensity_all << std::endl;
            }
        }

        Fileout::output_Etz (Parameter::Etz , nt, grid_2d->E_n0[0]);
        Fileout::fileout_1d_array(Parameter::intensity, intensity_array, nr_out, 1.);
        Fileout::fileout_1d_array(Parameter::intensity_r0, intensity_array, 1, 1.);
        
        Fileout::fileout_1d_array(Parameter::THzIntensity, intensityTHz_array, nr_out, 1.);
        Fileout::fileout_1d_array(Parameter::THzIntensity_r0, intensityTHz_array, 1, 1.);
        
        Fileout::fileout_1d_array(Parameter::plasmaDensity, plasmaDensity_array, nr_out, 1.);
        Fileout::fileout_1d_array(Parameter::plasmaDensity_r0, plasmaDensity_array, 1, 1.);


        Fileout::fileout_1d_array(Parameter::out_phase1, phase1_array, nr_out, 1.);
        Fileout::fileout_1d_array(Parameter::out_phase1_r0, phase1_array, 1, 1.);
        if(E_peak2 !=0){
            Fileout::fileout_1d_array(Parameter::out_phase2, phase2_array, nr_out, 1.);
            Fileout::fileout_1d_array(Parameter::out_phase2_r0, phase2_array, 1, 1.);

            Fileout::outputErzaDual(Parameter::Erza, nr_out, grid_2d->E_n0fft[index_omega1], grid_2d->E_n0fft[index_omega2]);
            Fileout::outputErzaDual(Parameter::Erza_r0, 0,   grid_2d->E_n0fft[index_omega1], grid_2d->E_n0fft[index_omega2]);
            Fileout::output_Erza(Parameter::Erza1, nr_out, grid_2d->E_n0fft[index_omega1]);
            Fileout::output_Erza(Parameter::Erza2, nr_out, grid_2d->E_n0fft[index_omega2]);
            Fileout::output_Erza(Parameter::Erza1_r0, 1,   grid_2d->E_n0fft[index_omega1]);
            Fileout::output_Erza(Parameter::Erza2_r0, 0,   grid_2d->E_n0fft[index_omega2]);
        }
        else if(E_peak2 == 0){
            Fileout::output_Erza(Parameter::Erza, nr_out, grid_2d->E_n0fft[index_omega1]);
            Fileout::output_Erza(Parameter::Erza_r0, 1, grid_2d->E_n0fft[index_omega1]);
        }
        
         
        //! solve--------------------------------------------------
        
        // complex A1_array[nr];
        // complex B1_array[nr];
        complex X1[nr];
        for (int num_t = 0; num_t < fft_num; num_t++)
        {
            
            CN(A0, A1[num_t], A2,
               B0, B1[num_t], B2, 
               grid_2d->E_n0fft[num_t], grid_2d->G_n0fft[num_t], X1);

            for (int num_r = 0; num_r < nr; num_r++)
            {
                grid_2d->E_n1fft[num_r][num_t] = X1[num_r];
                // 对num_t行数据求解，将得到的解放在E_n1fft 的第num_r行 num_t列
            }
        }

        //! ifft-------------------------------------------------
        for (int num_r = 0; num_r < nr; num_r++)
        {
            Fourier::ifft_c2r(nt, grid_2d->E_n1fft[num_r], grid_2d->E_n0[num_r]); 
            // 对r行数据进行ifft
        }
    }

    //Release Space----------------------------------------
    delete grid_2d;

    delete[] ionization_rate;
    delete[] electron_density;
    delete[] nonlinear_J;
    delete[] nonlinear_P;

    delete[] A0;
    delete[] A2;
    delete[] B0;
    delete[] B2;

    delete[] fft_outc_E;
	delete[] fft_outc_J;
    delete[] fft_outc_P;

    for(int i = 0; i < fft_num; i++){
        delete[] A1[i];
        delete[] B1[i];
    }
    delete[] A1;
    delete[] B1;
}